#!/usr/bin/env python3

"""
Given an array nums of n integers, are there elements a, b, c in nums such that a + b + c = 0?
Find all unique triplets in the array which gives the sum of zero
"""


class Solution:
    def three_sum(self, nums):
        val_to_index = {}
        nums.sort()
        for index, val in enumerate(nums):
            if val in val_to_index:
                val_to_index[val].append(index)
            else:
                val_to_index[val] = [index]
        result = []
        last = None
        for index, num in enumerate(nums):
            if num == last:
                # repeat, continue
                continue
            last = num
            remain = -num
            last_j_val = None
            for j_index in range(index + 1, len(nums)):
                j_val = nums[j_index]
                if j_val == last_j_val:
                    continue
                last_j_val = j_val
                j_remain = remain - j_val
                if j_remain in val_to_index:
                    for j_remain_index in val_to_index[j_remain]:
                        if j_remain_index > j_index:
                            result.append([num, j_val, j_remain])
                            break
        return result
    
if __name__ == '__main__':
    solution = Solution()
    result = solution.three_sum([0, 0, 0, 0])
    print(str(result))
    result = solution.three_sum([-1, 0, 1, 2, -1, -4])
